package com.ruoyi.common.sens;

import lombok.Data;

import java.util.*;

@Data
public class Trie {

    @Data
    public static class TireNode {
        private char value;
        private boolean end;
        private String word;
        private Map<Character, TireNode> children;


        public void initChildren() {
            if (this.children == null) {
                this.children = new HashMap<>();
            }
        }

        public TireNode createChild(char value) {
            TireNode node = new TireNode();
            node.setValue(value);
            this.initChildren();
            this.getChildren().put(value, node);
            return node;
        }
    }

    private TireNode root = new TireNode();

    private int size = 0;


    public void initWith(Collection<String> words) {
        for (String word : words) {
            TireNode currentNode = root;
            for(int index = 0; index < word.length(); index ++ ) {
                char ch = word.charAt(index);
                currentNode.initChildren();
                TireNode node = currentNode.getChildren().get(ch);
                if (node == null) {
                    node = currentNode.createChild(ch);
                    node.setEnd(false);
                }
                if (index == word.length() - 1) {
                    node.setEnd(true);
                    node.setWord(word);
                    this.size ++;
                }
                currentNode = node;
            }
        }
    }

    @Data
    public static class FindResult {
        int beginIndex;
        private String word;
    }

    public List<FindResult> find(String text) {
        return find(text, 0);
    }

    public List<FindResult> find(String text, int beginIndex) {
        List<FindResult> results = new ArrayList<>();

        for(int index = beginIndex; index < text.length() ; index++){
            TireNode currentNode = root;
            int i = index;
            while (true) {
                if (i >= text.length()) {
                    break;
                }
                char ch = text.charAt(i++);
                if (currentNode.getChildren() == null) {
                    break;
                }
                TireNode node = currentNode.getChildren().get(ch);
                if (node == null) {
                    break;
                }
                currentNode = node;
                if (currentNode.isEnd()) {
                    FindResult result = new FindResult();
                    result.setWord(node.getWord());
                    result.setBeginIndex(index);
                    results.add(result);
                    break;
                }
            }
            index = i - 1;
        }
        return results;
    }

    public String replaceWith(String text, List<FindResult> results) {
        return replaceWith(text, results, '*');
    }

    public String replaceWith(String text, List<FindResult> results, char mask) {
        StringBuilder maskText = new StringBuilder(text);
        for (FindResult result : results) {
            maskText.replace(
                result.getBeginIndex(),
                result.getBeginIndex() + result.getWord().length(),
                repeat(mask, result.getWord().length())
            );
        }
        return maskText.toString();
    }

    private static String repeat(char mask, int count)  {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < count; i++) {
            sb.append(mask);
        }
        return sb.toString();
    }

}
